import trimesh
from plyfile import PlyData, PlyElement
import numpy as np

def save_mesh2ply(mesh, save_path):
    points = mesh.vertices
    f = mesh.faces
    points = [(points[i,0], points[i,1], points[i,2]) for i in range(points.shape[0])]
    vertices = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4')])
    
    faces_building = []
    for i in range(0, f.shape[0]):
        faces_building.append(((f[i, :].tolist(),)))
    faces = np.array(faces_building, dtype=[("vertex_indices", "i4", (3,))])

    vertices = PlyElement.describe(vertices, 'vertex')
    faces = PlyElement.describe(faces, 'face')
    ply_data = PlyData([vertices, faces])
    ply_data.text = True
    ply_data.write(save_path)
    
def my_create_gripper_marker(color=[0, 0, 255], tube_radius=0.01, sections=6):
    cfl = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[
            [-7.27595772e-12, -4.10000000e-02, 6.59999996e-02],
            [-7.27595772e-12, -4.10000000e-02, 1.12169998e-01],
        ],
    )
    cfr = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[
            [-7.27595772e-12, 4.10000000e-02, 6.59999996e-02],
            [-7.27595772e-12, 4.10000000e-02, 1.12169998e-01],
        ],
    )
    cb1 = trimesh.creation.cylinder(
        radius=tube_radius, sections=sections, 
        segment=[[0, 0, 0], [0, 0, 6.59999996e-02]],
    )
    cb2 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, -4.100000e-02, 6.59999996e-02], [0, 4.100000e-02, 6.59999996e-02]],
    )

    tmp = trimesh.util.concatenate([cb1, cb2, cfr, cfl])
    tmp.visual.face_colors = color

    return tmp

def my_create_gripper_marker_jaco(color=[0, 0, 255], tube_radius=0.01, sections=6):
    cfl = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[
            [-7.27595772e-12, -4.10000000e-02, 6.59999996e-02],
            [-7.27595772e-12, -4.10000000e-02, 1.12169998e-01],
        ],
    )
    cfr = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[
            [-7.27595772e-12, 4.10000000e-02, 6.59999996e-02],
            [-7.27595772e-12, 4.10000000e-02, 1.12169998e-01],
        ],
    )
    cb1 = trimesh.creation.cylinder(
        radius=tube_radius, sections=sections, 
        segment=[[0, 0, 0], [0, 0, 6.59999996e-02]],
    )
    cb2 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, -4.100000e-02, 6.59999996e-02], [0, 4.100000e-02, 6.59999996e-02]],
    )

    tmp = trimesh.util.concatenate([cb1, cb2, cfr, cfl])
    tmp.visual.face_colors = color

    return tmp
    
if __name__ == '__main__':
    gripper_mesh = my_create_gripper_marker()
    save_mesh2ply(gripper_mesh, 'gripper.ply')
